/**
 * AES unit test: CBC encrypt/decrypt with 256-bit key
 **/

#include <stdio.h>

#include "pmsis.h"
#include "pmsis/drivers/aes.h"

//#define PRINTF(...)
#define PRINTF(...) printf(__VA_ARGS__)

#define DATA_SIZE 1024

static struct pi_device aes_device;

static uint8_t aes_key[32] = {
    0x00, 0x01, 0x02, 0x03,
    0x04, 0x05, 0x06, 0x07,
    0x08, 0x09, 0x0A, 0x0B,
    0x0C, 0x0D, 0x0E, 0x0F,
    0x10, 0x21, 0x32, 0x43,
    0x14, 0x25, 0x36, 0x47,
    0x18, 0x29, 0x3A, 0x4B,
    0x1C, 0x2D, 0x3E, 0x4F,
};

static uint8_t aes_iv[16] = {
    0x00, 0x11, 0x22, 0x33,
    0x04, 0x15, 0x26, 0x37,
    0x08, 0x19, 0x2A, 0x3B,
    0x0C, 0x1D, 0x2E, 0x3F,
};

static uint8_t src[DATA_SIZE] = {0};

static uint8_t dst[DATA_SIZE] = {0};
static uint8_t dst2[DATA_SIZE] = {0};

/* result from OpenSSL encryption */
static uint8_t exp_enc[DATA_SIZE] = {
    0xD6, 0x0D, 0xF3, 0xA3,
    0xF6, 0x43, 0x21, 0x73,
    0x62, 0x46, 0x5E, 0xC0,
    0xA6, 0x1C, 0x39, 0xD4,
    0xA5, 0xBB, 0x39, 0xC2,
    0xA6, 0x2B, 0x83, 0x41,
    0xAF, 0x92, 0x54, 0xF8,
    0x5F, 0x96, 0x4F, 0x41,
    0x31, 0x99, 0x23, 0x25,
    0x9B, 0xE7, 0x34, 0x03,
    0xA8, 0x43, 0xDA, 0x71,
    0x73, 0x96, 0x6C, 0xEF,
    0x62, 0x31, 0xE5, 0x3F,
    0xF9, 0xC1, 0x1C, 0xB9,
    0xA0, 0x60, 0x48, 0x09,
    0xA9, 0xBA, 0x38, 0x57,
    0x5C, 0x8A, 0x83, 0x9C,
    0xBE, 0x2B, 0xB1, 0xB6,
    0x40, 0x5A, 0x60, 0xE4,
    0x69, 0xBB, 0xB4, 0xD2,
    0x17, 0xFF, 0xE2, 0x3D,
    0x02, 0xFE, 0x34, 0xAA,
    0xF9, 0x15, 0x2A, 0x84,
    0x9C, 0x73, 0x53, 0x7F,
    0x7F, 0x0F, 0xDD, 0x5D,
    0x6A, 0x59, 0x92, 0x17,
    0x8D, 0x2E, 0x57, 0x2E,
    0x64, 0xBC, 0xD7, 0x03,
    0xA1, 0x82, 0xD0, 0xE9,
    0x2E, 0xE8, 0x66, 0x3B,
    0xA4, 0x8C, 0xCA, 0xE5,
    0xB9, 0x27, 0x30, 0x00,
    0x35, 0x92, 0x58, 0x03,
    0xFF, 0x54, 0x4A, 0xFD,
    0x3B, 0x2B, 0xEA, 0x84,
    0x4B, 0x31, 0x90, 0xE4,
    0x86, 0x6E, 0x42, 0x5E,
    0xD9, 0xFD, 0x77, 0xA8,
    0xC8, 0x4F, 0xB2, 0x32,
    0x6A, 0xA3, 0xA2, 0xF4,
    0x49, 0x03, 0x4B, 0x7C,
    0x32, 0xFF, 0x68, 0xC1,
    0x2D, 0x4F, 0x30, 0x93,
    0xF7, 0x1B, 0x28, 0xE6,
    0xBF, 0x73, 0xBE, 0x03,
    0x3B, 0x67, 0x62, 0x8E,
    0x7A, 0x7B, 0xC7, 0xDA,
    0x11, 0x3F, 0x2C, 0x1C,
    0x4A, 0xB6, 0x4A, 0xCC,
    0x75, 0xE1, 0x6B, 0x20,
    0x28, 0x1A, 0x4B, 0x1F,
    0x6C, 0x7A, 0xA0, 0x8D,
    0x07, 0xF2, 0xC2, 0xE5,
    0x8F, 0x56, 0xEE, 0x8C,
    0x60, 0x52, 0xE9, 0xA9,
    0x71, 0xA8, 0x3A, 0x84,
    0x08, 0xA5, 0x98, 0x32,
    0xDA, 0x85, 0xEE, 0xA1,
    0xD4, 0x94, 0x30, 0x1B,
    0x48, 0xB1, 0xD2, 0x3E,
    0x29, 0xF8, 0x12, 0xEA,
    0x2E, 0xC3, 0x31, 0xA0,
    0xE8, 0x2C, 0x06, 0x87,
    0x62, 0xB9, 0xC1, 0xA5,
    0x15, 0xDB, 0x6B, 0xFB,
    0xFD, 0x7F, 0x91, 0x29,
    0x10, 0x5D, 0xA0, 0x05,
    0xDC, 0x4F, 0xB7, 0x33,
    0xCA, 0x4B, 0x34, 0xE2,
    0x79, 0x80, 0x10, 0x72,
    0x10, 0x5D, 0x3A, 0x5A,
    0x97, 0x7C, 0xC4, 0x0B,
    0x7E, 0xD2, 0x51, 0x2B,
    0xD6, 0xDB, 0x09, 0x0F,
    0xC3, 0xA3, 0xA0, 0x50,
    0xE8, 0x3A, 0x3D, 0xA9,
    0x5B, 0xA2, 0x97, 0x3F,
    0x92, 0xBB, 0x01, 0xB0,
    0x04, 0xBE, 0x67, 0x96,
    0x5B, 0x8F, 0x13, 0x4C,
    0x88, 0xAB, 0x74, 0xC8,
    0x38, 0xB7, 0xA9, 0xC6,
    0x85, 0x6E, 0xD7, 0x3B,
    0xC9, 0x85, 0x5C, 0x8F,
    0x34, 0xC3, 0xE3, 0x99,
    0xDF, 0x10, 0x4D, 0x9D,
    0x49, 0x37, 0x67, 0x47,
    0x8C, 0x6A, 0xA7, 0xB3,
    0xE5, 0xA5, 0xB7, 0xC2,
    0x15, 0xBA, 0xBD, 0x7A,
    0x01, 0xC7, 0xE2, 0xF2,
    0x02, 0x6A, 0x4D, 0xE6,
    0x06, 0xD5, 0xBB, 0xEE,
    0x6E, 0x39, 0xD6, 0xD2,
    0x2E, 0x3C, 0x6F, 0xAB,
    0xBA, 0x15, 0xA1, 0x56,
    0x83, 0x6C, 0x33, 0x51,
    0x92, 0x86, 0xC0, 0x47,
    0x38, 0x74, 0x39, 0xFB,
    0x98, 0x98, 0x25, 0x15,
    0x07, 0x13, 0xB7, 0x4E,
    0xF5, 0x69, 0x87, 0x29,
    0x90, 0x58, 0x19, 0x90,
    0x6C, 0xD6, 0xAB, 0x95,
    0x29, 0xE7, 0x58, 0x6E,
    0xC1, 0x7B, 0x93, 0xCE,
    0xDF, 0x26, 0x07, 0x4E,
    0x34, 0xF4, 0x25, 0x93,
    0x34, 0x20, 0xC2, 0x61,
    0xA5, 0xEF, 0x6F, 0x04,
    0x97, 0x6E, 0x74, 0x03,
    0xD7, 0x8F, 0x18, 0xB2,
    0xC7, 0x53, 0xA6, 0xCF,
    0x25, 0x44, 0x78, 0x24,
    0xEB, 0x1F, 0x1C, 0x93,
    0xFA, 0xD6, 0xE2, 0x1C,
    0x89, 0x84, 0xD5, 0xC4,
    0x24, 0x13, 0xF2, 0x9F,
    0xEB, 0x5D, 0x53, 0x86,
    0xC1, 0x56, 0x02, 0x9E,
    0x2A, 0xA8, 0x3F, 0x01,
    0x36, 0xFB, 0xB4, 0x8C,
    0x0D, 0x57, 0x4C, 0x88,
    0x4D, 0x04, 0xE1, 0x4D,
    0xB5, 0x4C, 0xC1, 0x90,
    0x68, 0x3F, 0x00, 0x76,
    0xD1, 0x2B, 0x88, 0xE9,
    0x0E, 0xAB, 0x56, 0xA1,
    0x88, 0x96, 0x3F, 0x2E,
    0x8E, 0xD8, 0x1B, 0xE4,
    0x77, 0x54, 0x87, 0x2C,
    0x30, 0xDA, 0xC6, 0x6A,
    0x1D, 0xD5, 0x68, 0xDB,
    0x68, 0xE0, 0x0F, 0x35,
    0x75, 0x08, 0x6E, 0x44,
    0x94, 0xEB, 0x38, 0xC1,
    0x6E, 0x91, 0xB8, 0x24,
    0xEE, 0xF3, 0xD7, 0x52,
    0x59, 0xB2, 0xBB, 0xCC,
    0x26, 0xAB, 0xCF, 0x32,
    0x63, 0xE3, 0xD8, 0x9D,
    0x92, 0xAD, 0x39, 0xF6,
    0xF1, 0x20, 0xED, 0x21,
    0x3F, 0xBC, 0xEF, 0x9D,
    0xFA, 0xCB, 0xAD, 0xD3,
    0x5C, 0x40, 0x54, 0x45,
    0x7A, 0xCE, 0x17, 0xFE,
    0x21, 0x0D, 0x08, 0x41,
    0xD7, 0xC2, 0x53, 0xD6,
    0x53, 0x8B, 0x95, 0x0F,
    0x3A, 0x6F, 0xF1, 0x20,
    0x75, 0xA5, 0x4F, 0xED,
    0xB5, 0xE8, 0xD0, 0x5C,
    0x0F, 0x4A, 0xC0, 0x34,
    0xEF, 0xDF, 0x9F, 0x24,
    0x4F, 0x6D, 0x80, 0xF1,
    0x44, 0x28, 0x9F, 0x47,
    0xBA, 0x3B, 0x59, 0x60,
    0xBA, 0xB7, 0xF1, 0x84,
    0xC1, 0x54, 0x97, 0xCA,
    0x45, 0xF1, 0x61, 0x4F,
    0x95, 0xDD, 0xCC, 0x27,
    0x6D, 0x02, 0xE8, 0x9C,
    0xD4, 0xE5, 0xDD, 0xCB,
    0x65, 0x04, 0x13, 0x7D,
    0x97, 0xF6, 0xAD, 0x60,
    0xCC, 0xF2, 0xF0, 0xC5,
    0xC4, 0x40, 0x3E, 0x76,
    0x45, 0x74, 0xB1, 0xF8,
    0xC8, 0x4C, 0xDD, 0x95,
    0x6F, 0x55, 0xF7, 0x93,
    0x91, 0x96, 0x65, 0x93,
    0x2F, 0xD0, 0xD5, 0xEC,
    0x7D, 0x20, 0x49, 0x41,
    0x4E, 0x8F, 0x98, 0xA8,
    0x7C, 0x83, 0xB4, 0x10,
    0xFD, 0x56, 0x9D, 0x41,
    0xD8, 0x9E, 0xEA, 0xCC,
    0x5C, 0xF4, 0xAA, 0x11,
    0xF7, 0xDD, 0x53, 0xD5,
    0x8B, 0x8D, 0x18, 0x9F,
    0x45, 0xCA, 0x9B, 0x54,
    0xA8, 0x06, 0xC6, 0x01,
    0x7D, 0x13, 0x7A, 0x11,
    0x37, 0x61, 0x71, 0xF6,
    0xD5, 0x1E, 0xF5, 0x6F,
    0x56, 0x6C, 0xBA, 0xC7,
    0xE9, 0x17, 0x2F, 0xEB,
    0xF0, 0x59, 0x43, 0x4E,
    0x9B, 0xF9, 0xD2, 0x71,
    0xF2, 0xD9, 0x79, 0x0F,
    0xCE, 0x07, 0xD5, 0x5D,
    0xFF, 0xB5, 0xA6, 0x91,
    0x81, 0x78, 0xF4, 0xB6,
    0x1F, 0x59, 0xA1, 0xC7,
    0xA1, 0xB4, 0x09, 0x2A,
    0xDC, 0x4F, 0x42, 0x4B,
    0xF8, 0x88, 0xDE, 0x11,
    0xC9, 0xAA, 0x40, 0xBD,
    0xD3, 0xDD, 0xE2, 0x73,
    0x58, 0x3F, 0x9E, 0xFA,
    0x90, 0x59, 0xF3, 0xDD,
    0xC8, 0x96, 0xBF, 0xFC,
    0xE1, 0xB1, 0x57, 0x0E,
    0x1A, 0xFD, 0x9A, 0xD9,
    0xA2, 0x34, 0x1A, 0x6E,
    0x3D, 0x0F, 0x2D, 0x4A,
    0x45, 0xBC, 0x48, 0x3F,
    0xCD, 0x1C, 0x13, 0x6D,
    0x31, 0x99, 0xB8, 0xCA,
    0x89, 0x17, 0x13, 0xA7,
    0xBD, 0x3D, 0x07, 0x9D,
    0xFE, 0x00, 0xA1, 0xDF,
    0xF8, 0xD6, 0xBC, 0x5C,
    0x00, 0x80, 0x40, 0xDD,
    0x12, 0x09, 0x63, 0x5F,
    0x11, 0xD3, 0xC2, 0x0B,
    0x28, 0x38, 0x4E, 0x69,
    0x5B, 0x23, 0x6B, 0xCB,
    0x0F, 0x95, 0x30, 0x7D,
    0x4A, 0x00, 0xDA, 0x65,
    0x5E, 0x8F, 0x54, 0x00,
    0x7D, 0xBC, 0xA4, 0xD9,
    0x24, 0x79, 0x4D, 0x14,
    0x06, 0x6D, 0x76, 0xC7,
    0x7F, 0x35, 0xFA, 0x08,
    0xE1, 0x59, 0x6B, 0xD8,
    0x85, 0x81, 0x74, 0x27,
    0xDE, 0x2B, 0xC5, 0xD7,
    0x8E, 0x46, 0xF5, 0x7A,
    0x10, 0x58, 0xCA, 0x1B,
    0xA0, 0x58, 0x8F, 0x0C,
    0xAC, 0xB9, 0xBC, 0xFA,
    0xEF, 0x22, 0x75, 0x10,
    0x01, 0xBF, 0xC2, 0x18,
    0x50, 0x2B, 0x34, 0xF5,
    0x42, 0xC7, 0x6C, 0xE2,
    0xA8, 0x6C, 0x5B, 0x20,
    0xB7, 0x6C, 0x55, 0x72,
    0x33, 0x3F, 0x64, 0xB5,
    0xB9, 0xAF, 0xA5, 0x92,
    0xAF, 0xF7, 0x30, 0x1B,
    0x15, 0xC0, 0x36, 0x11,
    0x84, 0x3D, 0xD8, 0x17,
    0x69, 0x4E, 0x93, 0xF9,
    0x7A, 0x0F, 0x2F, 0xE0,
    0x9B, 0x6D, 0xE5, 0x88,
    0xFE, 0x0D, 0xA3, 0x84,
    0x31, 0x2E, 0x4B, 0xD6,
    0xA9, 0x46, 0x4C, 0xA2,
    0xE8, 0xFF, 0xB6, 0x18,
    0xAA, 0x2F, 0x14, 0x4D,
    0x8E, 0xD5, 0x66, 0x14,
    0xA2, 0xEC, 0xB6, 0xAC,
    0xA8, 0x98, 0x8E, 0x97,
    0xBC, 0xD0, 0x3B, 0xD9
};

static int aes_open(void)
{
    struct pi_aes_conf conf;
    pi_aes_conf_init(&conf);
    conf.mode = PI_AES_MODE_CBC;
    conf.key_len = PI_AES_KEY_256;
    conf.key = (uint32_t*) aes_key;
    conf.iv = (uint32_t*) aes_iv;

    pi_open_from_conf(&aes_device, &conf);
    int ret = pi_aes_open(&aes_device);

    if(!ret)
    {
        PRINTF("AES open succeeded\n");
    }
    else
    {
        PRINTF("AES open failed\n");
    }

    return ret;
}

static int aes_close(void)
{
    pi_aes_close(&aes_device);
    return 0;
}

static void data_init(void)
{
    for (int i = 0; i < DATA_SIZE; i++)
    {
        src[i] = i & 0xFF;
    }
}

static int aes_crypt(void)
{
    int ret = 0;
    int errors = 0;

    pi_udma_fifo_conf_t conf = {0};
    pi_device_t fifo_input = {0};
    pi_device_t fifo_output = {0};
    pi_task_t block = {0};

    pi_task_block(&block);

    data_init();

    /* initialize fifos */
    pi_udma_fifo_conf_init(&conf);
    pi_open_from_conf(&fifo_input, &conf);
    pi_open_from_conf(&fifo_output, &conf);

    ret = pi_udma_fifo_open(&fifo_input);
    if (0 != ret)
    {
        printf("Error while opening input fifo\n");
        pmsis_exit(-1);
    }
    ret = pi_udma_fifo_open(&fifo_output);
    if (0 != ret)
    {
        printf("Error while opening output fifo\n");
        pmsis_exit(-1);
    }

    /* set buffer sizes */
    uint32_t size = DATA_SIZE;
    pi_udma_fifo_ioctl(&fifo_input, PI_UDMA_FIFO_SET_SIZE, &size);
    pi_udma_fifo_ioctl(&fifo_output, PI_UDMA_FIFO_SET_SIZE, &size);

    /* get fifo ids */
    uint32_t fifo_input_id, fifo_output_id;
    pi_udma_fifo_ioctl(&fifo_input, PI_UDMA_FIFO_GET_ID, &fifo_input_id);
    pi_udma_fifo_ioctl(&fifo_output, PI_UDMA_FIFO_GET_ID, &fifo_output_id);

    /* processing */

    printf("Encrypt data\n");
    ret = pi_aes_encrypt_async(&aes_device, (uint32_t*) fifo_input_id, (uint32_t*) fifo_output_id, DATA_SIZE / sizeof(uint32_t), &block);
    printf("- Pushing (long)\n");
    for(int i = 0; i < DATA_SIZE; i++)
    {
        pi_udma_fifo_ioctl(&fifo_input, PI_UDMA_FIFO_PUSH8, &(src[i]));
    }
    printf("- Popping (long)\n");
    for(int i = 0; i < DATA_SIZE; i++)
    {
        pi_udma_fifo_ioctl(&fifo_output, PI_UDMA_FIFO_POP8, &(dst[i]));
    }
    printf("- Verifying\n");
    for (int i = 0; i < DATA_SIZE; i++)
    {
        if (dst[i] != exp_enc[i])
        {
            errors++;
            printf("ERROR: encrypted data invalid - dst: %02x, exp_enc: %02x\n", dst[i], exp_enc[i]);
        }
    }
    pi_aes_ioctl(&aes_device, PI_AES_STOP_FIFO_MODE, NULL);

    printf("\nDecrypting data\n");
    ret += pi_aes_decrypt_async(&aes_device, (uint32_t*) fifo_input_id, (uint32_t*) fifo_output_id, DATA_SIZE / sizeof(uint32_t), &block);
    printf("- Pushing (long)\n");
    for(int i = 0; i < DATA_SIZE; i++)
    {
        pi_udma_fifo_ioctl(&fifo_input, PI_UDMA_FIFO_PUSH8, &(dst[i]));
    }
    printf("- Popping (long)\n");
    for(int i = 0; i < DATA_SIZE; i++)
    {
        pi_udma_fifo_ioctl(&fifo_output, PI_UDMA_FIFO_POP8, &(dst2[i]));
    }
    printf("- Verifying\n");
    for (int i = 0; i < DATA_SIZE; i++)
    {
        if (src[i] != dst2[i])
        {
            errors++;
            printf("ERROR: decrypted data invalid - expected %02x, got %02x\n", src[i], dst2[i]);
        }
    }

    pi_aes_ioctl(&aes_device, PI_AES_STOP_FIFO_MODE, NULL);

    pi_udma_fifo_close(&fifo_input);
    pi_udma_fifo_close(&fifo_output);

    printf("\n");
    return ret + errors;
}

/**
 * \brief launch the main test
 *
 * \warning hangs indefinitely
 */
int test_main(void)
{
    PRINTF("Application start\n");

    int ret = 0;

    ret = aes_open();

    ret += aes_crypt();

    ret += aes_close();

    if(ret)
    {
        printf("test returned with %d errors\n", ret);
    }
    else
    {
        printf("Test succeeded\n");
    }
    pmsis_exit(ret);
    while(1);
    return 0;
}

int main(void)
{
    return pmsis_kickoff((void*) test_main);
}

